#!/usr/bin/env python3
import time
from rich.progress import track

from utils.global_utils import *
from utils.ev_utils import *
from utils.converter import *
from utils.visualization import *
from event_ros.msg import CameraData, Float64Mat


data_dir: str = os.path.join(config["data_dir"].replace("~", os.environ["HOME"]), "camera")

# pub node
nodePub: rospy.Publisher = None


if __name__ == "__main__":
    rospy.init_node("cam_replay", anonymous = True)

    nodePub = rospy.Publisher("/event_ros/sub_event_res", CameraData, queue_size = 1)
    publisherInitialized = True

    replay_config = config["replay"]
    outer_sleep = lambda: time.sleep(replay_config["countdown"])
    inner_sleep = lambda: time.sleep(replay_config["real_step_ms"] / 1000.0) # ms -> s
    duration_base = replay_config["sim_step_ms"] * 1000.0 # ms -> us

    sub_dirs = [sub_dir for sub_dir, dirs, files in list(os.walk(data_dir)) if "events.npy" in files]
    total = len(sub_dirs)
    for idx, sub_dir in enumerate(sub_dirs):
        if rospy.is_shutdown():
            break

        print("Now playing %s (%d/%d)." % (sub_dir.split(os.sep)[-1], idx + 1, total))
        events = np.load(os.path.join(sub_dir, "events.npy"))
        events[:, 0] *= 1e6 # s -> us

        
        pos_events = events.copy()
        pos_events[:, 1] = 1
        pos_events[:, 2] = np.clip(pos_events[:, 2] + np.random.randint(-2, 2, pos_events[:, 2].shape), 0, 799)
        pos_events[:, 3] = np.clip(pos_events[:, 3] + np.random.randint(-2, 2, pos_events[:, 3].shape), 0, 1279)
        neg_events = events.copy()
        neg_events[:, 1] = 0
        events = np.concatenate([pos_events, neg_events], axis = 0)

        events = events[events[:, 3] > (-0.1 * events[:, 2] + 540)]

        if os.path.isfile(os.path.join(sub_dir, "imu.npy")):
            imus = np.load(os.path.join(sub_dir, "imu.npy"))
        else:
            imus = np.zeros((0, 7), dtype = np.float64)

        duration = 0.0
        last_duration = 0.0
        earliest, latest = np.min(events[:, 0]), np.max(events[:, 0])
        events[:, 0] -= earliest
        while duration <= latest - earliest:
            if rospy.is_shutdown():
                break
            event_slice = events[(events[:, 0] >= last_duration) & (events[:, 0] < duration)]
            if imus.shape[0]:
                imu_slice = imus[(imus[:, 0] >= last_duration) & (imus[:, 0] < duration)]
            else:
                imu_slice = np.zeros((0, 7), dtype = np.float64)
            last_duration = duration
            time_now = get_timestamp()
            secs_now, nsecs_now = int(secs(time_now)), int(nsecs(time_now))
            frame_id = json.dumps({"path": sub_dir, "env": True})
            if event_slice.shape[0]:
                data_msg = CameraData()
                data_msg.events = events_4col_to_5col(event_slice).flatten().astype(np.uint32).tolist()
                data_msg.imus = imu_slice.flatten().astype(np.float64).tolist()
                data_msg.header.stamp.secs = secs_now
                data_msg.header.stamp.nsecs = nsecs_now
                data_msg.header.frame_id = frame_id
                data_msg.eventcol = ["t_sec", "t_nsec", "p", "y", "x"]
                data_msg.imucol = ["t", "angvel_z", "angvel_y", "angvel_x", "accel_z", "accel_y", "accel_x"]
                nodePub.publish(data_msg)
            inner_sleep()

            # 改为固定时间
            duration += duration_base
        
        time_now = get_timestamp()
        secs_now, nsecs_now = int(secs(time_now)), int(nsecs(time_now))
        frame_id = json.dumps({"path": sub_dir, "env": None})
        data_msg = CameraData()
        data_msg.header.stamp.secs = secs_now
        data_msg.header.stamp.nsecs = nsecs_now
        data_msg.header.frame_id = frame_id
        nodePub.publish(data_msg)
    
    print("Playing done.")
    while not rospy.is_shutdown():
        continue